{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use LLamaIndexQueryEngine to query Markdown files \n",
    "\n",
    "This notebook demonstrates the use of the `LLamaIndexqueryEngine` for retrieval-augmented question answering over documents. It shows how to set up the engine with Docling parsed Markdown files, and execute natural language queries against the indexed data. \n",
    "\n",
    "The `LlamaIndexQueryEngine` provides an efficient way to query vectorDBs using any LlamaIndex's [vector store](https://docs.llamaindex.ai/en/stable/module_guides/storing/vector_stores/).\n",
    "\n",
    "We use some Markdown (.md) files as input, feel free to try your own text or Markdown documents.\n",
    "\n",
    "You can create and add this ChromaDBQueryEngine to [DocAgent](https://docs.ag2.ai/latest/docs/user-guide/reference-agents/docagent) to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-vector-stores-chroma==0.4.1\n",
    "%pip install llama-index==0.12.16\n",
    "%pip install llama-index llama-index-vector-stores-pinecone==0.4.4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load LLM configuration\n",
    "\n",
    "This demonstration requires an `OPENAI_API_KEY` to be in your environment variables. See [our documentation](https://docs.ag2.ai/latest/docs/user-guide/advanced-concepts/llm-configuration-deep-dive) for guidance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import autogen\n",
    "\n",
    "config_list = autogen.config_list_from_json(env_or_file=\"../OAI_CONFIG_LIST\")\n",
    "\n",
    "assert len(config_list) > 0\n",
    "print(\"models to use: \", [config_list[i][\"model\"] for i in range(len(config_list))])\n",
    "\n",
    "# Put the OpenAI API key into the environment\n",
    "os.environ[\"OPENAI_API_KEY\"] = config_list[0][\"api_key\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### In the first example, we build a LLamaIndexQueryEngine instance using ChromaDB. \n",
    "Refer to this [link](https://docs.trychroma.com/production/containers/docker) for running Chromadb in a Docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chromadb import HttpClient\n",
    "from llama_index.vector_stores.chroma import ChromaVectorStore\n",
    "\n",
    "# we need to set up LlmaIndex's ChromaVectorStore\n",
    "# Refer to https://docs.llamaindex.ai/en/stable/examples/vector_stores/chroma_metadata_filter/ for more information\n",
    "chroma_client = HttpClient(\n",
    "    host=\"host.docker.internal\",\n",
    "    port=8000,\n",
    ")\n",
    "\n",
    "# Option 1: get an existing collection\n",
    "# use get_collection to get an existing collection\n",
    "chroma_collection = chroma_client.get_collection(\"default_collection\")\n",
    "\n",
    "# Option 2: create a new collection\n",
    "# chroma_collection = chroma_client.create_collection(\"default_collection\")\n",
    "\n",
    "# Create the Chroma vector store\n",
    "chroma_vector_store = ChromaVectorStore(chroma_collection=chroma_collection)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can use the LlamaIndex `chroma_vector_store` to create our AG2 `LLamaIndexQueryEngine` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "from autogen.agentchat.contrib.rag import LlamaIndexQueryEngine\n",
    "\n",
    "chroma_query_engine = LlamaIndexQueryEngine(\n",
    "    vector_store=chroma_vector_store,\n",
    "    llm=OpenAI(model=\"gpt-4.1\", temperature=0.0),  # Default model for querying, change if needed\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize the database with input docs and query it with the engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_dir = (\n",
    "    \"/workspaces/ag2/test/agents/experimental/document_agent/pdf_parsed/\"  # Update to match your input directory\n",
    ")\n",
    "input_docs = [input_dir + \"nvidia_10k_2024.md\"]  # Update to match your input documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Option 1: initialize the database and add new documents\n",
    "chroma_query_engine.init_db(new_doc_paths_or_urls=input_docs)\n",
    "\n",
    "# Option 2: connect to the database without initializing it\n",
    "# chroma_query_engine.connect_db()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# question = \"How much money did Nvidia spend in research and development\"\n",
    "question = \"What was the latest quarter's GAAP revenue?\"\n",
    "answer = chroma_query_engine.query(question)\n",
    "print(answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great, we got the data we needed. Now, let's add another document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_docs = [input_dir + \"Toast_financial_report.md\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chroma_query_engine.add_docs(new_doc_paths_or_urls=new_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And query again from the same database but this time for another corporate entity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"How much money did Toast earn in 2024?\"\n",
    "answer = chroma_query_engine.query(question)\n",
    "print(answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pinecone\n",
    "\n",
    "### In the second example, we build a similar LLamaIndexQueryEngine instance, but on top of Pinecone.\n",
    "Refer to https://docs.llamaindex.ai/en/stable/examples/vector_stores/PineconeIndexDemo/ for more details on how to set up Pinecone and  PineconeVectorStore\n",
    "\n",
    "Please put your Pinecone API key in an environment variable called `PINECONE_API_KEY`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pinecone import Pinecone, ServerlessSpec\n",
    "\n",
    "# Load the Pinecode API key and create the Pinecone object\n",
    "api_key = os.environ[\"PINECONE_API_KEY\"]\n",
    "pc = Pinecone(api_key=api_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dimensions are for text-embedding-ada-002, which PineconeVectorStore uses for embedding text by default\n",
    "\n",
    "# Create an index named ag2\n",
    "pc.create_index(\n",
    "    name=\"ag2\",\n",
    "    dimension=1536,\n",
    "    metric=\"euclidean\",\n",
    "    spec=ServerlessSpec(cloud=\"aws\", region=\"us-east-1\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.vector_stores.pinecone import PineconeVectorStore\n",
    "\n",
    "# Create the vector store\n",
    "pinecone_index = pc.Index(\"ag2\")\n",
    "pinecone_vector_store = PineconeVectorStore(pinecone_index=pinecone_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pinecone_query_engine = LlamaIndexQueryEngine(\n",
    "    vector_store=pinecone_vector_store,\n",
    "    llm=OpenAI(model=\"gpt-4.1\", temperature=0.0),  # Default model for querying, change if needed\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the database and add new documents\n",
    "pinecone_query_engine.init_db(new_doc_paths_or_urls=input_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Query the Pinecone query engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"How much money did Nvidia spend in research and development\"\n",
    "answer = pinecone_query_engine.query(question)\n",
    "print(answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add another document"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_docs = [input_dir + \"Toast_financial_report.md\"]\n",
    "pinecone_query_engine.add_docs(new_doc_paths_or_urls=new_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Query again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"How much money did Toast earn in 2024?\"\n",
    "answer = pinecone_query_engine.query(question)\n",
    "print(answer)"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Use any LlamaIndex vector store as a Query Engine",
   "tags": [
    "agents",
    "documents",
    "RAG",
    "docagent",
    "chroma",
    "chromadb",
    "pinecone"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
